import os
import numpy as np
import nibabel as nib
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt
import skimage
from scipy import ndimage
import torchvision.transforms as transforms
import random

class MRIDataset(Dataset):
    def __init__(self, root_dir, modalities=('axi', 'sag', 'cor'), max_slices=3, 
                crop_size=128, target_size=128, augment=None, downsample_ratio=0.5):
        self.root_dir = root_dir
        self.modalities = modalities
        self.max_slices = max_slices
        self.crop_size = crop_size
        self.augment = augment
        self.target_size = target_size
        self.downsample_ratio = downsample_ratio
        
        # 加载标签
        self.labels = pd.read_csv(os.path.join(root_dir, 'label.csv'),
                        header=0,
                        names=['p_id', 'label']).set_index('p_id')
        self.labels['label'] = self.labels['label'].astype(int) - 1  # 转换为0-based
        assert self.labels['label'].between(0, 1).all(), "标签必须在0-1范围内"
        
        # 存储所有通过完整性检查的患者ID
        self.all_patient_ids = [] 
        self.class0_ids = [] 
        self.class1_ids = [] 
        for pid in self.labels.index:
            if self._check_data_complete(pid):
                self.all_patient_ids.append(pid)
                label = self.labels.loc[pid, 'label']
                if label == 0:
                    self.class0_ids.append(pid)
                else:
                    self.class1_ids.append(pid)

        # 初始化当前患者ID列表（这个列表会在每个epoch被重置）
        self.current_patient_ids = []
        self.reset_epoch()  # 初始化当前列表
        self.print_class_counts()

    def reset_epoch(self):
        """在每个epoch开始时调用，重新采样第二类样本以形成当前epoch的列表"""
        # 第一类全部保留
        class0_sampled = self.class0_ids
        # 第二类根据下采样比例随机采样
        num_class1 = int(len(self.class1_ids) * self.downsample_ratio)
        num_class1 = max(1, num_class1)  # 至少一个
        if num_class1 < len(self.class1_ids):
            class1_sampled = random.sample(self.class1_ids, num_class1)
        else:
            class1_sampled = self.class1_ids

        self.current_patient_ids = class0_sampled + class1_sampled
        random.shuffle(self.current_patient_ids)  # 打乱顺序
        class_counts = {0: len(class0_sampled), 1: len(class1_sampled)}
        print(f"当前epoch的类别分布: 第0类={class_counts[0]}, 第1类={class_counts[1]}, 总计={len(self.current_patient_ids)}")

    def print_class_counts(self):
        class0_total = len(self.class0_ids)
        class1_total = len(self.class1_ids)
        print(f"总样本分布: 第0类={class0_total}, 第1类={class1_total}, 总计={class0_total+class1_total}")

    def _check_data_complete(self, pid):
        for mod in self.modalities:
            img_path = os.path.join(self.root_dir, 'img', f'{pid}_T2_{mod}_000.nii.gz')
            roi_path = os.path.join(self.root_dir, 'roi', f'{pid}_T2_{mod}_roi.nii.gz')
            
            if not (os.path.exists(img_path) and os.path.exists(roi_path)):
                return False
            try:
                img = nib.load(img_path)
                roi = nib.load(roi_path)
            except Exception as e:
                print(f"加载文件失败：{img_path} 或 {roi_path}, 错误：{e}")
                return False
            img_shape = img.shape[:2]
            roi_shape = roi.shape[:2]
            if img_shape != roi_shape:
                print(f"患者 {pid} 的 {mod} 模态图像尺寸 {img_shape} 与 ROI 尺寸 {roi_shape} 不匹配")
                return False
            max_allowed_ratio = 1.2
            h, w = img_shape
            if h == 0 or w == 0:
                print(f"患者 {pid} 的 {mod} 模态存在零尺寸：{img_shape}")
                return False
            current_ratio = max(h / w, w / h)
            if current_ratio > max_allowed_ratio:
                print(f"患者 {pid} 的 {mod} 模态长宽比过大：{img_shape} (比例 {current_ratio:.2f})")
                return False
            
        return True

    def _crop_around_center(self, image, center_y, center_x, is_roi=False):
        """根据给定中心坐标进行裁剪"""
        target_size = self.crop_size
        half_size = target_size // 2
        
        # 计算理论裁剪范围
        y_start = center_y - half_size
        y_end = center_y + half_size
        x_start = center_x - half_size
        x_end = center_x + half_size
        
        # 计算实际可裁剪范围
        actual_y_start = max(0, y_start)
        actual_y_end = min(image.shape[0], y_end)
        actual_x_start = max(0, x_start)
        actual_x_end = min(image.shape[1], x_end)
        
        # 执行实际裁剪
        cropped = image[actual_y_start:actual_y_end,
                        actual_x_start:actual_x_end]
        
        # 计算需要填充的边界
        pad_y_before = max(0, -y_start)
        pad_y_after = max(0, y_end - image.shape[0])
        pad_x_before = max(0, -x_start)
        pad_x_after = max(0, x_end - image.shape[1])
        
        # 使用边缘像素填充
        padded = np.pad(cropped,
                        ((pad_y_before, pad_y_after),
                        (pad_x_before, pad_x_after)),
                        mode='edge')
        
        img_resize = skimage.transform.resize(padded, (self.target_size, self.target_size))
        if is_roi:
            img_resize = skimage.transform.resize(padded, (self.target_size, self.target_size))
            binary_img = (img_resize > 0.5).astype(np.uint8)
            inverted_img = 1 - binary_img
            dist_transform = ndimage.distance_transform_edt(inverted_img)
            expansion_radius = 4
            expansion_mask = (dist_transform <= expansion_radius) & (inverted_img == 1)
            expanded_roi = np.logical_or(binary_img, expansion_mask).astype(np.float32)
            img_resize = expanded_roi

        return img_resize.astype(np.float32)
    
    def _apply_augmentation(self, img, roi):
        if not self.augment:
            return img, roi
        img_aug = img.copy()
        roi_aug = roi.copy()
        
        # 1. 随机旋转 (角度偏转)
        if random.random() < 1:
            angle = random.uniform(-15, 15)
            img_aug = skimage.transform.rotate(
                img_aug, angle, resize=False, mode='edge', order=1)
            # 最近邻插值以保持二值特性
            roi_aug = skimage.transform.rotate(
                roi_aug, angle, resize=False, mode='edge', order=0)
        
        # 2. 随机高斯噪声
        if random.random() < 0.2: 
            noise = np.random.normal(0, 0.05, img_aug.shape).astype(np.float32)
            img_aug = img_aug + noise
            
            img_aug = np.clip(img_aug, np.min(img), np.max(img))
        if random.random() < 0.2:  
            # 随机亮度调整 (-10%到+10%)
            brightness_factor = random.uniform(0.9, 1.1)
            img_aug = img_aug * brightness_factor
            contrast_factor = random.uniform(0.8, 1.2)
            mean_val = np.mean(img_aug)
            img_aug = (img_aug - mean_val) * contrast_factor + mean_val
            
            img_aug = np.clip(img_aug, np.min(img), np.max(img))
        
        if random.random() < 0.2:  # 50%概率应用遮挡
            max_occlusions = 3  # 最大遮挡数量
            occlusion_size_range = (0.1, 0.25)  # 遮挡尺寸比例范围

            h, w = img_aug.shape
            for _ in range(random.randint(1, max_occlusions)):
                # 生成随机遮挡尺寸
                occ_height = int(h * random.uniform(*occlusion_size_range))
                occ_width = int(w * random.uniform(*occlusion_size_range))
                
                # 生成随机位置
                x = random.randint(0, h - occ_height)
                y = random.randint(0, w - occ_width)
                
                # 应用矩形遮挡（填充0）
                img_aug[x:x+occ_height, y:y+occ_width] = 0

        return img_aug, roi_aug

    def __len__(self):
        return len(self.current_patient_ids) 

    def __getitem__(self, idx):
        pid = self.current_patient_ids[idx]
        label = self.labels.loc[pid, 'label']
        
        multimodal_img = []
        multimodal_roi = []

        for mod in self.modalities:
            img_path = os.path.join(self.root_dir, 'img', f'{pid}_T2_{mod}_000.nii.gz')
            roi_path = os.path.join(self.root_dir, 'roi', f'{pid}_T2_{mod}_roi.nii.gz')
            
            img = nib.load(img_path)
            roi = nib.load(roi_path)
            
            img_data = img.get_fdata().astype(np.float32)  # (H, W, D)
            roi_data = roi.get_fdata().astype(np.float32)
            
            # 1. Find slice with largest ROI
            max_roi_area = -1
            max_slice_idx = -1
            for slice_idx in range(roi_data.shape[2]):
                current_area = np.sum(roi_data[..., slice_idx] > 0)
                if current_area > max_roi_area:
                    max_roi_area = current_area
                    max_slice_idx = slice_idx
            
            if max_slice_idx == -1:
                max_slice_idx = img_data.shape[2] // 2
                print(f"患者 {pid} 的 {mod} 模态没有有效的ROI，使用中间层切片")

            if max_slice_idx == img_data.shape[2]:
                max_slice_idx == img_data.shape[2] - 2
                print(f"患者{pid} 的 最大ROI图位于最后一张切片，使用前一张切片为中心")
            
            if max_slice_idx == 0:
                max_slice_idx == 1
                print(f"患者{pid} 的 最大ROI图位于第一张切片，使用后一张切片为中心")

            roi_max_slice = roi_data[..., max_slice_idx]
            coords = np.argwhere(roi_max_slice > 0)
            center_y = img_data.shape[0] // 2
            center_x = img_data.shape[1] // 2
            if len(coords) > 0:
                min_y, min_x = coords.min(axis=0)
                max_y, max_x = coords.max(axis=0)
                center_y = (min_y + max_y) // 2
                center_x = (min_x + max_x) // 2
            
            # 3. Get 3 slices (center ±1)
            selected_indices = []
            for offset in range(-1, 2):
                selected_idx = max_slice_idx + offset
                selected_indices.append(selected_idx)

            img_slices_tensor = []  # 存储每个切片的张量
            roi_slices_tensor = []   # 存储每个ROI切片的张量

            for slice_idx in selected_indices:
                # 取出对应位置的切片，nibabel数据的H*W
                img_slice = img_data[..., slice_idx]
                roi_slice = roi_data[..., slice_idx]
                
                img_cropped = self._crop_around_center(img_slice, center_y, center_x, is_roi=False)
                roi_cropped = self._crop_around_center(roi_slice, center_y, center_x, is_roi=True)

                img_cropped = (img_cropped - np.min(img_cropped)) / (np.max(img_cropped) - np.min(img_cropped))
                img_cropped = (img_cropped - 0.5) / 0.5

                img_slice, roi_slice = self._apply_augmentation(img_cropped, roi_cropped)

                img_slice_tensor = torch.tensor(img_slice, dtype=torch.float32)
                roi_slice_tensor = torch.tensor(roi_slice, dtype=torch.float32)
                
                img_slices_tensor.append(img_slice_tensor)
                roi_slices_tensor.append(roi_slice_tensor)#用一个列表接住img_slice_tensor

            multimodal_img.append(torch.stack(img_slices_tensor))  # (3, H, W)
            multimodal_roi.append(torch.stack(roi_slices_tensor)) # (3, H, W)

        multimodal_img = torch.stack(multimodal_img)
        multimodal_roi = torch.stack(multimodal_roi)

        img_output = multimodal_img  # (M, D, H, W)，[:,0]代表取每个位面的第一张切片，然后有三个位面
        roi_output = multimodal_roi  # 0代表输入第一个位面的所有切片[0, :]，

        return (
            img_output,  # 取第1个位面的第2张切片[1:2,:,:]
            roi_output,
            torch.tensor(label, dtype=torch.long),
            # pid
        )

def visualize_attention(img_batch, roi_batch, labels):
    # 获取首个样本数据
    img, roi, label = img_batch[0].numpy(), roi_batch[0].numpy(), labels[0].item()
    mod_names = ['Axial', 'Sagittal', 'Coronal']
    
    # 创建图形窗口
    fig, axs = plt.subplots(3, 3, figsize=(15, 15))
    plt.suptitle(f"Label: {label} | Slice: 0/{img.shape[1]-1}", fontsize=16)
    
    # 初始化切片索引
    current_slice = [0]  # 使用列表实现闭包内的可变性
    max_slice = img.shape[1] - 1
    
    # 存储图像对象
    image_plots = [[None for _ in range(3)] for _ in range(3)]
    
    # 初始化显示
    for mod in range(3):
        # 原始图像
        image_plots[0][mod] = axs[0, mod].imshow(img[mod, 0], cmap='gray')
        axs[0, mod].set_title(f"{mod_names[mod]} Image")
        
        # ROI掩膜
        image_plots[1][mod] = axs[1, mod].imshow(roi[mod, 0], cmap='Reds', alpha=0.5)
        axs[1, mod].set_title(f"{mod_names[mod]} ROI Mask")
        
        # 叠加显示
        image_plots[2][mod] = axs[2, mod].imshow(img[mod, 0], cmap='gray')
        image_plots[2][mod] = axs[2, mod].imshow(roi[mod, 0], cmap='Reds', alpha=0.3)
        axs[2, mod].set_title("Overlay")
    
    def update_display():
        """更新所有模态的显示"""
        fig.suptitle(f"Label: {label} | Slice: {current_slice[0]}/{max_slice}", fontsize=16)
        for mod in range(3):
            # 更新原始图像
            image_plots[0][mod].set_data(img[mod, current_slice[0]])
            
            # 更新ROI掩膜
            image_plots[1][mod].set_data(roi[mod, current_slice[0]])
            
            # 更新叠加显示（需要重新创建以避免透明度混合问题）
            axs[2, mod].images[0].set_data(img[mod, current_slice[0]])
            axs[2, mod].images[1].set_data(roi[mod, current_slice[0]])
        
        fig.canvas.draw_idle()
    
    def on_scroll(event):
        """鼠标滚轮事件处理"""
        if event.button == 'up':
            current_slice[0] = min(current_slice[0] + 1, max_slice)
        elif event.button == 'down':
            current_slice[0] = max(current_slice[0] - 1, 0)
        update_display()
    
    # 绑定滚轮事件
    fig.canvas.mpl_connect('scroll_event', on_scroll)
    
    # 添加键盘事件支持
    def on_key(event):
        """键盘事件处理（左右方向键）"""
        if event.key == 'right':
            current_slice[0] = min(current_slice[0] + 1, max_slice)
        elif event.key == 'left':
            current_slice[0] = max(current_slice[0] - 1, 0)
        update_display()
    
    fig.canvas.mpl_connect('key_press_event', on_key)
    
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    dataset = MRIDataset(root_dir='D:/PyChrom/PythonProject/2.medical_image/DeepTrip/datasets/pT/three_label_data/train', 
                         augment=True,
                         downsample_ratio=0.5)
    print(len(dataset))
    loader = DataLoader(dataset, batch_size=4,shuffle=False)
    
    img_batch, roi_batch, labels = next(iter(loader))

    visualize_attention(img_batch, roi_batch, labels)

    print(f"IMG   shape: {img_batch.shape}")  # 应该输出 (B, M=3, D=3, H=224, W=224)
    print(f"ROI   shape: {roi_batch.shape}")
    # print(f"Batch shape: {mask_batch.shape}")  # 应该输出 (B, M=3, D=3, H=224, W=224)
